What does this PR do ?

Help mitigate the performance overheads from clearing cache which is introduced in #926 by making clear cache behaviour optional.

Issues

This PR resolves #1036.

Usage

This PR introduces an additional option for dtensor_config. To clear cache between microbatch iterations, add:

policy:
    dtensor_cfg:
        # .....
        clear_cache_every_n_steps: 1

A warning will be displayed before a training iteration indicating that clear cache has been turned on and its potential performance overheads.

(DTensorPolicyWorker[rank=0] pid=3016383) /lustre/fs1/portfolios/coreai/users/yubog/RL/nemo_rl/models/policy/dtensor_policy_worker.py:635: UserWarning: Emptying cache every 1 microbatches
, doing so unnnecessarily would incur a large performance overhead.
(DTensorPolicyWorker[rank=0] pid=3016383)   warnings.warn(f"Emptying cache every {empty_cache_steps} microbatches, doing so unnnecessarily would incur a large performance overhead.")

Below is the speed comparison of the default SFT experiment with clear cache every 1 microbatch vs. no clear cache:

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
• Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

• ...